/*
** $Id: sluac.c,v 1.0.0 2015/11/2 11:34:15 yinqiang
** Lua compiler (saves bytecodes to files encrypt by xxtea)
** Base on the Lua, Lua.org, PUC-Rio, Brazil (http://www.lua.org)
** 
************************************************************************************
** Copyright (c) 2015, yinqiang
** All rights reserved.
** 
** Redistribution and use in source and binary forms, with or without modification,
** are permitted provided that the following conditions are met:
** 
** * Redistributions of source code must retain the above copyright notice, this
**   list of conditions and the following disclaimer.
** 
** * Redistributions in binary form must reproduce the above copyright notice, this
**   list of conditions and the following disclaimer in the documentation and/or
**   other materials provided with the distribution.
** 
** THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
** ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
** WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
** DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
** ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
** (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
** LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
** ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
** (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
** SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
** 
************************************************************************************
** Copyright (C) 1994-2012 Lua.org, PUC-Rio.  All rights reserved.
**
** Permission is hereby granted, free of charge, to any person obtaining
** a copy of this software and associated documentation files (the
** "Software"), to deal in the Software without restriction, including
** without limitation the rights to use, copy, modify, merge, publish,
** distribute, sublicense, and/or sell copies of the Software, and to
** permit persons to whom the Software is furnished to do so, subject to
** the following conditions:
**
** The above copyright notice and this permission notice shall be
** included in all copies or substantial portions of the Software.
**
** THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
** EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
** MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
** IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
** CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
** TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
** SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
************************************************************************************
*/

#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>
#include <sys/stat.h>
#ifdef __APPLE__
  #include <sys/uio.h>
#else
  #include <io.h>
#endif

#include "lua.h"
#include "lauxlib.h"

#include "ldo.h"
#include "lfunc.h"
#include "lmem.h"
#include "lobject.h"
#include "lopcodes.h"
#include "lstring.h"
#include "lundump.h"

#include "base64/libbase64.h"
#include "xxtea/xxtea.h"

#define PROGNAME    "sluac"      /* default program name */
#define OUTPUT      PROGNAME ".out" /* default output file */

#define SLUAC_VERSION   "1.0.1"
#define SLUAC_COPYRIGHT "Copyright (c) 2015, yinqiang"

#define BUFF_APD_STEP 4*1024   /* xxtea buffer append step */

#ifndef FILEPATH_MAX
  #define FILEPATH_MAX 4096
#endif

struct fs_node_s
{
    int isdir;
    char path[FILEPATH_MAX];
    struct fs_node_s *next;
};

#define FS_ISDIR(N)   ((N)->isdir==1)

#define MAX(a,b)   ((a)>(b)?(a):(b))

static int quiet = 0;                   /* no printing */
static int stripping = 0;               /* strip debug information? */
static char Output[] = { OUTPUT };      /* default output file name */
static const char* input = NULL;
static const char* output = Output;     /* actual output file name */
static const char* progname = PROGNAME; /* actual program name */

static int xxteaing = 0;
static char xxtea_key[256] = {'\0'};
static char* xxtea_sig = "sluac";
static size_t xxtea_sig_len = 0;

static char* buff = NULL;
static size_t buff_len = 0;
static size_t buff_max = 0;

static void fatal(const char* message)
{
    fprintf(stderr, "%s: %s\n", progname, message);
    exit(EXIT_FAILURE);
}

static void cannot(const char* what, const char* name)
{
    fprintf(stderr, "%s: cannot %s %s: %s\n", progname, what, name, strerror(errno));
    exit(EXIT_FAILURE);
}

static void usage(const char* message)
{
    if (*message == '-')
        fprintf(stderr, "%s: unrecognized option " LUA_QS "\n", progname, message);
    else
        fprintf(stderr, "%s: %s\n", progname, message);
    fprintf(stderr,
        "usage: %s [options].\n"
        "Available options are:\n"
        "  -i  path  intput path of dir\n"
        "  -o  path  output path of dir\n"
        "  -x  key   xxtea key, base64 encoded\n"
        "  -xs sig   xxtea sig string\n"
        "  -s        strip debug information\n"
        "  -v        show version information\n"
        "  -q        no printing\n",
        progname);
    exit(EXIT_FAILURE);
}

#define IS(s)   (strcmp(argv[i],s)==0)

static int doargs(int argc, char* argv[])
{
    int i;
    int version = 0;
    
    if (argv[0] != NULL && *argv[0] != 0)
        progname = argv[0];

    for (i=1; i<argc; i++)
    {
        if (*argv[i] != '-')    /* end of options; keep it */
            break;
        else if (IS("-i"))      /* input file */
        {
            input = argv[++i];
            if (input == NULL || *input == 0)
                usage(LUA_QL("-i") " needs argument");
        }
        else if (IS("-o"))      /* output file */
        {
            output = argv[++i];
            if (output == NULL || *output == 0)
                usage(LUA_QL("-o") " needs argument");
        }
        else if (IS("-s"))      /* strip debug information */
            stripping = 1;
        else if (IS("-v"))      /* show version */
            ++version;
        else if (IS("-x"))      /* xxtea key */
        {
            xxteaing = 1;
            Base64decode((char*)xxtea_key, argv[++i]);
        }
        else if (IS("-xs"))     /* xxtea sig */
            xxtea_sig = argv[++i];
        else if (IS("-q"))
            ++quiet;
        else                    /* unknown option */
            usage(argv[i]);
    }
    if (version)
    {
        printf("%s  %s, %s\n", SLUAC_VERSION, LUA_COPYRIGHT, SLUAC_COPYRIGHT);
        if (version == argc-1)
            exit(EXIT_SUCCESS);
    }
    if (input == NULL || *input == 0)
        usage(LUA_QL("-i") " needs argument");
    return i;
}

#define toproto(L,i) (clvalue(L->top+(i))->l.p)

static inline int
is_pow_of_2(uint32_t x) {
    return !(x & (x-1));
}

static inline uint32_t
next_pow_of_2(uint32_t x) {
    if ( is_pow_of_2(x) )
        return x;
    x |= x>>1;
    x |= x>>2;
    x |= x>>4;
    x |= x>>8;
    x |= x>>16;
    return x+1;
}

static int writer(lua_State* L, const void* p, size_t size, void* u)
{
    UNUSED(L);
    if (buff_len + size > buff_max)
    {
        buff_max = next_pow_of_2(MAX(buff_max << 1, buff_max + size));
        buff = realloc(buff, buff_max);
        if (buff == NULL)
            fatal("not enough memory for buff");
    }
    memcpy(buff + buff_len, p, size);
    buff_len += size;
    return 0;
}

struct Smain {
    const char* dst;
    const char* src;
    const char* name;
};

static int pmain(lua_State* L)
{
    struct Smain* s = (struct Smain*)lua_touserdata(L, 1);
    const Proto* f;
    char* xxtea_bytes = NULL;
    size_t ret_len;
    
    if (luaL_loadfile(L, s->src) != 0)
        fatal(lua_tostring(L, -1));

    f = toproto(L, -1);

    FILE* D = fopen(s->dst, "wb");
    if (D == NULL)
        cannot("open", s->dst);

    buff_len = 0;
    memset(buff, 0, buff_max);

    lua_lock(L);
    luaU_dump(L, f, writer, D, stripping);
    lua_unlock(L);
    
    if (xxteaing)
    {
        ret_len = 0;
        xxtea_bytes = (char*)xxtea_encrypt(
            (unsigned char*)buff, buff_len,
            (unsigned char*)xxtea_key,  strlen((char*)xxtea_key),
            (xxtea_long *)&ret_len);
        
        if (xxtea_bytes == NULL || ret_len <= 0)
            cannot("encrypt", s->dst);
        else
        {
            fwrite(xxtea_sig, xxtea_sig_len, 1, D);
            fwrite(xxtea_bytes, ret_len, 1, D);
            free(xxtea_bytes);
            xxtea_bytes = NULL;
        }
    }
    else
        fwrite(buff, buff_len, 1, D);

    if (ferror(D))
        cannot("write", s->dst);
    if (fclose(D))
        cannot("close", s->dst);

    if (!quiet)
        fprintf(stdout,
            " > get bytes [%4ld KB] %s\n",
            MAX(1, (xxteaing ? ret_len : buff_len) / 1024),
            s->name);

    return 0;
}

static struct fs_node_s* fs_lsdir(const char* dir, struct fs_node_s** tail)
{
    struct dirent *entry;
    struct stat st;
    struct fs_node_s *head, *p, *tmp;
    char path[FILEPATH_MAX], root[FILEPATH_MAX];
    DIR *d = NULL;
    int n;

    if (dir == NULL)
        return NULL;

    n = strlen(dir);
    if (n <= 0 || n > FILEPATH_MAX)
        return NULL;
    sprintf(root, (dir[n-1] != '/' ? "%s/" : "%s"), dir);

    d = opendir(root);
    if (d == NULL)
        return NULL;

    p = malloc(sizeof(struct fs_node_s));
    sprintf(p->path, "%s", root);
    p->isdir = 1;
    p->next = NULL;
    head = p;

    entry = NULL;
    while ((entry = readdir(d)) != NULL)
    {
        if (entry->d_name[0] != '.')
        {
            sprintf(path, "%s%s", root, entry->d_name);
            if (stat(path, &st) == 0)
            {
                if (S_ISDIR(st.st_mode))
                {
                    tmp = NULL;
                    p->next = fs_lsdir(path, &tmp);
                    if (p->next == NULL)
                        cannot("list dir", path);
                    else
                        p = tmp;
                }
                else
                {
                    tmp = malloc(sizeof(struct fs_node_s));
                    sprintf(tmp->path, "%s", path);
                    tmp->isdir = 0;
                    tmp->next = NULL;
                    p->next = tmp;
                    p = tmp;
                }
            }
            else
                cannot("get stat", path);
        }
    }
    closedir(d);

    if (tail != NULL)
        *tail = p;

    return head;
}

static void fs_free(struct fs_node_s* head)
{
    struct fs_node_s *p, *tmp;

    for (p=head; p!=NULL; )
    {
        tmp = p;
        p = p->next;
        free(tmp);
    }
}

int main(int argc, char* argv[])
{
    lua_State* L;
    struct Smain s;
    struct fs_node_s *files, *p;
    char path[FILEPATH_MAX], root[FILEPATH_MAX];
    int offset, done = 1;

    doargs(argc, argv);
    
    L = lua_open();
    if (L == NULL)
        fatal("not enough memory for state");
    
    buff_max = BUFF_APD_STEP;
    buff = malloc(buff_max);
    if (buff == NULL)
        fatal("not enough memory for buff");

    if (xxteaing)
    {
        xxtea_sig_len = strlen(xxtea_sig);
    }

    sprintf(root, (*(output+strlen(output)-1) != '/' ? "%s/" : "%s"), output);
    offset = strlen(input);
    if (*(input+offset-1) != '/')
        ++offset;

    files = fs_lsdir(input, NULL);
    if (files == NULL)
        return EXIT_FAILURE;

    for (p=files; p!=NULL; p=p->next)
    {
        sprintf(path, "%s%s", root, p->path + offset);
        if (FS_ISDIR(p))
        {
#ifdef __APPLE__
            if (mkdir(path, 755) != 0)
#else
            if (mkdir(path) != 0)
#endif
                fprintf(stderr, "dir exisit, %s\n", path);
        }
        else
        {
            s.src = p->path;
            s.dst = path;
            s.name = p->path+offset;
            if (lua_cpcall(L, pmain, &s) != 0)
            {
                fatal(lua_tostring(L, -1));
                done = 0;
            }
        }
    }

    if (buff != NULL)
        free(buff);

    lua_close(L);
    fs_free(files);

    if (!done)
    {
        fatal("fatal");
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}
